// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::ffi::{c_char, c_int, CStr};

use tracing::{debug, error};

use cudax::{driver, runtime};

use crate::cache;

fn cuda_check(err: driver::CUresult, file: &str, line: u32) {
    if err != driver::cudaError_enum_CUDA_SUCCESS {
        let mut err_name: *const c_char = std::ptr::null();
        unsafe {
            driver::cuGetErrorName(err, &mut err_name);
        }

        let err_str = if err_name.is_null() {
            "Unknown error".to_string()
        } else {
            unsafe { CStr::from_ptr(err_name).to_string_lossy().into_owned() }
        };

        error!("CUDA error at {}:{}: {}", file, line, err_str);
    }
}

fn ensure_cuda_context(device_id: c_int) {
    unsafe {
        let mut current_ctx: driver::CUcontext = std::ptr::null_mut();
        cuda_check(driver::cuCtxGetCurrent(&mut current_ctx), file!(), line!());

        if current_ctx.is_null() {
            debug!(
                "No active context on this thread. Activating primary context for device {}.",
                device_id
            );

            let mut device: driver::CUdevice = 0;
            let mut primary_ctx: driver::CUcontext = std::ptr::null_mut();

            cuda_check(
                driver::cuDeviceGet(&mut device, device_id),
                file!(),
                line!(),
            );

            cuda_check(
                driver::cuDevicePrimaryCtxRetain(&mut primary_ctx, device),
                file!(),
                line!(),
            );

            cuda_check(driver::cuCtxSetCurrent(primary_ctx), file!(), line!());
        }
    }
    debug!("ensure_cuda_context ok");
}

pub fn load_function(module_name: &str, func_name: &CStr) -> Result<driver::CUfunction, u32> {
    debug!("[serverx] load_function({:?})...", func_name);

    let mut device_id = 0;

    let ret = unsafe { runtime::cudaGetDevice(&mut device_id) };
    if ret != 0 {
        error!("Cannot get device, ret={}", ret);
        return Err(ret);
    }

    // cuda_check(err, file!(), line!());
    ensure_cuda_context(device_id);

    cache::get_or_load_function(module_name, func_name)
}
